import os

import cv2
import numpy as np
from keras import Sequential
from keras_preprocessing.image import ImageDataGenerator

from tensorflow.keras.layers import Convolution2D, BatchNormalization, Activation, MaxPooling2D, Dropout, Flatten, Dense
import tensorflow as tf
import matplotlib.pyplot as plt
from tensorflow.python.keras.callbacks_v1 import TensorBoard

class Model(object):
    def build_model(self):
        self.model = Sequential()
        self.model.add(
            Convolution2D(
                filters=32,
                kernel_size=(5, 5),
                padding='same',
                input_shape=(150, 150, 3))
        )
        self.model.add(BatchNormalization())
        self.model.add(Activation('relu'))
        self.model.add(
            MaxPooling2D(
                pool_size=(2, 2),
                strides=(2, 2),
                padding='same')
        )


        self.model.add(Convolution2D(filters=64, kernel_size=(5, 5), padding='same'))
        self.model.add(BatchNormalization())
        self.model.add(Activation('relu'))
        self.model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same'))
        self.model.add(Dropout(0.15))


        self.model.add(Convolution2D(filters=64, kernel_size=(5, 5), padding='same'))
        self.model.add(BatchNormalization())
        self.model.add(Activation('relu'))
        self.model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same'))
        self.model.add(Dropout(0.15))



        return self.model

model = Model()
base_model = model.build_model()
base_model.load_weights("Pre-VGG.h5", by_name=True, skip_mismatch=True)
base_model.trainable = False

model = Sequential()

model.add(base_model)

model.add(Flatten())
model.add(Dense(512))
model.add(BatchNormalization())
model.add(Activation('relu'))
model.add(Dropout(0.5))

model.add(Dense(128))
model.add(BatchNormalization())
model.add(Activation('relu'))
model.add(Dropout(0.5))

model.add(Dense(2))
model.add(BatchNormalization())
model.add(Activation('softmax'))
model.summary()




model.compile(optimizer=tf.keras.optimizers.RMSprop(learning_rate=0.0005),
              loss=tf.keras.losses.CategoricalCrossentropy(from_logits=False),
              metrics=["accuracy"])

callbacks = [tf.keras.callbacks.ModelCheckpoint(filepath='Pre-VGG_trim.h5',
                                                save_best_only=True,
                                                save_weights_only=True,
                                                monitor='val_loss')]

log_filepath = r'logs'
tb_cb = TensorBoard(log_dir=log_filepath, write_images=1, histogram_freq=1)
cbks = [tb_cb]
# model.fit_generator(generate_one_epoch(train_X, train_y, batch_size=batch_size),
#                     steps_per_epoch=len(train_X) // batch_size,
#                     epochs=100, callbacks=cbks)

# tensorflow2.1 recommend to using fit
image_path = os.path.join("train")
train_dir = os.path.join(image_path, "train")
validation_dir = os.path.join(image_path, "val")
train_image_generator = ImageDataGenerator(rescale=1. / 255,
                                               horizontal_flip=True)
validation_image_generator = ImageDataGenerator(rescale=1. / 255)

batchsize = 64

train_data_gen = train_image_generator.flow_from_directory(directory=train_dir,
                                                           batch_size=batchsize,
                                                           shuffle=True,
                                                           target_size=(150, 150),
                                                           class_mode='categorical')
total_train = train_data_gen.n

val_data_gen = validation_image_generator.flow_from_directory(directory=validation_dir,
                                                              batch_size=batchsize,
                                                              shuffle=False,
                                                              target_size=(150, 150),
                                                              class_mode='categorical')
total_val = val_data_gen.n
print("using {} images for training, {} images for validation.".format(total_train,
                                                                       total_val))
# 训练次数
epochs=10

history = model.fit(x=train_data_gen,
                    steps_per_epoch=total_train // batchsize,
                    epochs=epochs,
                    validation_data=val_data_gen,
                    validation_steps=total_val // batchsize,
                    callbacks=[callbacks, cbks])

# plot loss and accuracy image
history_dict = history.history
train_loss = history_dict["loss"]
train_accuracy = history_dict["accuracy"]
val_loss = history_dict["val_loss"]
val_accuracy = history_dict["val_accuracy"]

# plot loss and accuracy image


# figure 1
plt.figure()
loss = plt.subplot(221)
loss.plot(range(epochs), train_loss, label='train_loss')
loss.plot(range(epochs), val_loss, label='val_loss')
plt.legend()
plt.xlabel('epochs')
plt.ylabel('loss')


accuracy = plt.subplot(222)
accuracy.plot(range(epochs), train_accuracy, label='train_accuracy')
accuracy.plot(range(epochs), val_accuracy, label='val_accuracy')
plt.legend()
plt.xlabel('epochs')
plt.ylabel('accuracy')
plt.savefig("TrainData.jpg")
plt.show()
